import React, { useState, useEffect, forwardRef, useImperativeHandle, useCallback } from 'react';
import { Tree as TreeAntd } from 'antd';

const { TreeNode } = TreeAntd;

function updateTreeData(list, key, children) {
  return list.map((node) => {
    if (node.key === key) {
      return { ...node, children };
    }
    if (node.children) {
      return { ...node, children: updateTreeData(node.children, key, children) };
    }
    return node;
  });
}

function deleteTreeData(list, key) {
  return list
    .filter(x => x.id !== key)
    .map((x) => {
      if (x.children) {
        return { ...x, children: deleteTreeData(x.children, key) };
      }
      return x;
    });
}

const Tree = (props, ref) => {
  const {
    treeData: initTreeData = [],
    getChildrenData,
    onSelect = () => { },
    defaultSelectedKeys = [],
    ...rest
  } = props;
  const [treeData, setTreeData] = useState(initTreeData);
  const [selected, setSelected] = useState(
    initTreeData.find(item => item.key === defaultSelectedKeys[0])
  );

  useEffect(() => {
    setTreeData(initTreeData);
  }, [initTreeData]);

  function onLoadData(data) {
    const { key, children, props: propsData } = data;
    return new Promise((resolve) => {
      if (children) {
        resolve();
        return;
      }
      getChildrenData(propsData.data, (dataNew) => {
        if (dataNew.length > 0) {
          setTreeData(origin =>
            updateTreeData(origin, key, dataNew),
          );
          resolve();
        }
      });
    });
  }

  // 删除树节点
  const deleteTreeNode = useCallback((record) => {
    setTreeData(origin =>
      deleteTreeData(origin, record.id),
    );
  }, []);

  useImperativeHandle(ref, () => ({
    deleteTreeNode,
    // 用于在树外部获取当前选中的节点 dataRef 属性
    getSelected: () => selected,
  }), [deleteTreeNode, selected]);

  const handleSelect = (selectedKeys, e) => {
    setSelected(e?.node?.props?.data);
    onSelect(selectedKeys, e, treeData);
  };

  const renderTreeNodes = data =>
    data.map((item) => {
      if (item.children) {
        return (
          <TreeNode title={item.title} key={item.key} dataRef={item} {...item}>
            {renderTreeNodes(item.children)}
          </TreeNode>
        );
      }
      return <TreeNode key={item.key} {...item} isLeaf={!item.hasChild} />;
    });

  return (
    <TreeAntd loadData={onLoadData} onSelect={handleSelect} defaultSelectedKeys={defaultSelectedKeys} {...rest}>
      {renderTreeNodes(treeData)}
    </TreeAntd>
  );
};

export default forwardRef(Tree);
